import json
from fastapi import APIRouter, Query
from fastapi.responses import JSONResponse
from typing import Optional
from langchain_community.llms import Ollama
from langchain_core.prompts import PromptTemplate

router = APIRouter()

# Ollama LLM 实例
ollama = Ollama(base_url='http://172.16.21.38:11436', model="qwen3:0.6b")

@router.get("/name")
async def generate_name(
    county: str = Query(..., description="起名字的风格或国家"),
    boy: str = Query(..., description="男孩常用名字示例"),
    girl: str = Query(..., description="女孩常用名字示例"),
    history: Optional[str] = Query(None, description="可选历史消息 JSON 字符串")):
    try:
        messages = json.loads(history) if history else []
        # 构造提示消息
        prompt = f"你是一个起名大师,请模仿示例起3个{county}名字,比如男孩经常被叫做{boy},女孩经常被叫做{girl}"
        messages.append({"role": "user", "content": prompt})
    except Exception as e:
        return JSONResponse(content={"success": False, "error": f"解析 history 失败: {str(e)}"})

    try:
        # 调用 Ollama 的 invoke 方法
        response = ollama.invoke(messages)
        return JSONResponse(content={"success": True, "response": response})
    except Exception as e:
        return JSONResponse(content={"success": False, "error": str(e)})


@router.get("/named2")
async def generate_name(
    county: str = Query(..., description="起名字的风格或国家"),
    boy: str = Query(..., description="男孩常用名字示例"),
    girl: str = Query(..., description="女孩常用名字示例"),
    history: Optional[str] = Query(None, description="可选历史消息 JSON 字符串")):
    try:
        messages = json.loads(history) if history else []
        # 构造提示消息
        # prompt = f"你是一个起名大师,请模仿示例起3个{county}名字,
                    #           比如男孩经常被叫做{boy},女孩经常被叫做{girl}"
        template = PromptTemplate.from_template(
            "你是一个起名大师,请模仿示例起3个具有{county}特色的名字,"
            "示例：男孩常用名{boy},女孩常用名{girl}。"
            "请返回以逗号分隔的列表形式。仅返回逗号分隔的列表，不要返回其他内容。"
        )
        prompt_str = template.format(county=county, boy=boy, girl=girl)
        messages.append({"role": "user", "content": prompt_str})
    except Exception as e:
        return JSONResponse(content={"success": False, "error": f"解析 history 失败: {str(e)}"})

    try:
        # 调用 Ollama 的 invoke 方法
        response = ollama.invoke(messages)
        return JSONResponse(content={"success": True, "response": response})
    except Exception as e:
        return JSONResponse(content={"success": False, "error": str(e)})

# GET http://127.0.0.1:8000/named2?county=美国男孩&boy=Sam&girl=Lucy

@router.get("/test")
async def test():
    return JSONResponse(content={"message": "Hello World"})
